#!/usr/bin/env python
# -*- coding: utf-8 -*-
import numpy as np
from hyperopt import fmin, hp, tpe
from typing import Dict, Union, List
from utils import get_qm_energy, generate_lammps_input_slab, generate_lammps_input_bin, get_lammps_energy, run_lammps
from utils import generate_lammps_input_in, get_qm_force, get_lammps_force, generate_lammps_pairwise, get_lammps_energy_mol
#from utils import generate_lammps_pairwise_part
#from utils import generate_lammps_pairwise_part_v1
#from utils import generate_lammps_pairwise_part_m1
#from utils import generate_lammps_pairwise_part_m2
#from utils import generate_lammps_pairwise_part_m3
#from utils import generate_lammps_pairwise_part_m4
#from utils import generate_lammps_pairwise_part_m5
from utils import generate_lammps_pairwise_part_v2
from sklearn.metrics import mean_squared_error, r2_score
import os, sys
import re
import math
import subprocess
from lammps import PyLammps
from utils import generate_datafiles, xdatcar2xyz
#from vasp2lmpdatav1 import xdatcar2xyzv1
#from vasp2lmpdatav2 import xdatcar2xyzv2
#from vasp2lmpdatav3 import xdatcar2xyzv3
from utils import generate_lammps_input_bin_2
from utils import generate_lammps_input_bin_3
from utils import generate_lammps_input_bin_4

def main():
    E_qm = []
    F_qm = []
    MolQM = []
    MolMM = []
    slab_qm = -966.6195
    l1=list(range(0,30))
    l2=[1111]
    Alll=[i for i in l1 if i not in l2]
    Activelll = []
    Activelli = []
    for i in Alll:
        Activelll.append(i)
        Activelli.append(Alll.index(i))
        E_qm.append((get_qm_energy('/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/VASP-%d/OUTCAR' % i) - get_qm_energy('/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/MOLE-%d/OUTCAR' % i) - slab_qm))
        #F_qm.append(get_qm_force('/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/VASP-%d/OUTCAR' % i))
        MolMM.append(get_lammps_energy_mol('/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/MOLE-%d/molEne%d' % (i, i)))
        MolQM.append(get_qm_energy('/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/MOLE-%d/OUTCAR' % i))
    
    E_target = [j * 23.06 for j in E_qm]
    E_target_scale = [j / 1.0 for j in E_target]
    F_target = F_qm

    op = open('FIT-ana30.log%d' % int(sys.argv[1]), 'w')

    def objective(params: Dict[str, float]):
        E_lammps = []
        F_lammps = []
        k=-1
        for i in Alll:
            k += 1
            j=1+int(i/6)
            if i <31:
              generate_lammps_input_slab('/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/fitting/DataFile/ana-slab.data', 'slab.data', params)
              subprocess.run('bash changeCHG-slab.sh %f' % params['chg'] ,shell=True, cwd='/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/fitting')
              run_lammps('ana-slab.in', 'slab')
            Natoms = generate_datafiles('/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/fitting/DataFile/ana_%d.data' % j )
            if i <6:
              xdatcar2xyz('/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/VASP-%d/CONTCAR' % i, '/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/fitting/DataFile/ana_%d.data-part2' % j , 'newpart2')
            elif i>=6 and i<12:
              xdatcar2xyz('/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/VASP-%d/CONTCAR' % i, '/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/fitting/DataFile/ana_%d.data-part2' % j , 'newpart2')
            elif i>=12 and i<18:
              xdatcar2xyz('/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/VASP-%d/CONTCAR' % i, '/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/fitting/DataFile/ana_%d.data-part2' % j , 'newpart2')
            elif i>=18 and i<24:
              xdatcar2xyz('/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/VASP-%d/CONTCAR' % i, '/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/fitting/DataFile/ana_%d.data-part2' % j , 'newpart2')
            elif i>=24 and i<30:
              xdatcar2xyz('/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/VASP-%d/CONTCAR' % i, '/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/fitting/DataFile/ana_%d.data-part2' % j , 'newpart2')
            else:
              xdatcar2xyz('/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/VASP-%d/CONTCAR' % i, '/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/fitting/DataFile/ana_%d.data-part2' % j , 'newpart2')
            os.system('cat /share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/fitting/DataFile/ana_%d.data-part1 newpart2 /share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/fitting/DataFile/ana_%d.data-part3 > paa-%d.data' %( j, j, i))
            if i>=6 and i<12:
              generate_lammps_input_bin_2('paa-%d.data' % i, 'bin-%d.data' % i, params)
            elif i>=12 and i<18:
              generate_lammps_input_bin_3('paa-%d.data' % i, 'bin-%d.data' % i, params)
            elif i>=18 and i<24:
              generate_lammps_input_bin_4('paa-%d.data' % i, 'bin-%d.data' % i, params)
            elif i>=24 and i<30:
              generate_lammps_input_bin('paa-%d.data' % i, 'bin-%d.data' % i, params)
            else:
              generate_lammps_input_bin('paa-%d.data' % i, 'bin-%d.data' % i, params)
            #generate_lammps_input_bin('paa-%d.data' % i, 'bin-%d.data' % i, params)
            subprocess.run('bash changeCHG-bin.sh %f bin-%d.data bin-%dC.data' % (params['chg'], i, i),shell=True, cwd='/share/workspace/zhaolingci/task-BASF-Jun2021/Nov-2021/anatase/fitting')
            #generate_lammps_input_in('ana-paa.in', 'paa-%d.in' % i, 'bin-%dC.data' % i, 'min-%d.xyz' % i)
#consider pairwise
            generate_lammps_input_in('ana-%d.in' % j, 'paa-%d.in' % i, 'bin-%dC.data' % i, 'min-%d.xyz' % i)
            ##generate_lammps_input_in('ana-m%d.in' % j, 'paa-%d.in' % i, 'bin-%dC.data' % i, 'min-%d.xyz' % i)
            #generate_lammps_pairwise_part_v1(params)
            generate_lammps_pairwise_part_v2(params)
            #generate_lammps_pairwise_part_m1(params)
            #generate_lammps_pairwise_part_m2(params)
            #generate_lammps_pairwise_part_m3(params)
            #generate_lammps_pairwise_part_m4(params)
            #generate_lammps_pairwise_part_m5(params)
            ##generate_lammps_pairwise_part(params)
            ##generate_lammps_pairwise(params)
            run_lammps('paa-%d.in' % i, 'binary')
            #energy = get_lammps_energy('binary') - get_lammps_energy('slab') - MolMM[i]
            energy = get_lammps_energy('binary') - get_lammps_energy('slab') - MolMM[k]
            E_lammps.append(energy)
            #F_lammps.append(get_lammps_force('forces'))
        op.write(' E_MSE:   ')
        op.write(str(mean_squared_error(E_target, E_lammps)))
        op.write(' E_R_square:   ')
        op.write(str(r2_score(E_target, E_lammps)))
        op.write(' E_RMSE:   ')
        op.write(str(math.sqrt(mean_squared_error(E_target, E_lammps))))
        op.write(' lmp_ads_ene:   ')
        op.write(str(E_lammps))
        op.write(' params:   ')
        op.write(str(params)+ '\n')
        op.write('#########'+'\n')
        #op.write(' F_MSE:   ')
        #op.write(str(mean_squared_error(F_target, F_lammps)))
        #op.write(' F_R_square:   ')
        #op.write(str(r2_score(F_target, F_lammps)))
        #op.write(' lmp_ads_force:   ')
        #op.write(str(F_lammps))
        op.write(' params:   ')
        op.write(str(params)+ '\n')
        op.write('#########'+'\n')
        #print(F_lammps)
        return mean_squared_error(E_target, E_lammps)
        #return -r2_score(E_target, E_lammps)
        #return mean_squared_error(F_target, F_lammps) + mean_squared_error(E_target, E_lammps)
        #E_lammps_scale = [ j / 1.0 for j in E_lammps]
        #return mean_squared_error(F_target, F_lammps) + mean_squared_error(E_target_scale, E_lammps_scale)

    SPACE = {
#'chg': 0.65, 'eps_1': 0.33, 'eps_2': 0.39, 'eps_3': 0.308, 'eps_4': 0.142, 'eps_o1ti5': 0.25, 'sigma_1': 3.0, 'sigma_2': 3.0, 'sigma_3': 4.15, 'sigma_4': 3.65, 'sigma_o1ti5': 2.45,
#        'sigma_c1o2': hp.quniform('sigma_c1o2', low=1.5, high=5.0, q=0.02),
#        'eps_c1o2': hp.quniform('eps_c1o2', low=0.01, high=0.95, q=0.002)
# 'chg': 0.65, 'eps_1': 0.45, 'eps_2': 0.31, 'eps_3': 0.104, 'eps_4': 0.18, 'eps_o1o2': 0.556, 'eps_o1ti5': 0.362, 'sigma_1': 2.7, 'sigma_2': 3.1, 'sigma_3': 3.95, 'sigma_4': 4.3, 'sigma_o1o2': 3.1, 'sigma_o1ti5': 2.45,
#        'sigma_Ho2': hp.quniform('sigma_Ho2', low=0.9, high=5.5, q=0.05),
#        'eps_Ho2': hp.quniform('eps_Ho2', low=0.01, high=0.95, q=0.002)
#'eps_1': 0.04, 'eps_2': 0.16, 'eps_3': 0.04, 'eps_4': 0.606, 'eps_ho2': 0.624, 'eps_o1ti5': 0.232, 'sigma_1': 2.30, 'sigma_2': 2.5, 'sigma_3': 3.25, 'sigma_4': 4.65, 'sigma_ho2': 1.55, 'sigma_o1ti5': 2.40,
#        'chg': hp.quniform('chg', low=0.4, high=0.9, q=0.005)
        'chg': 0.65,
        'sigma_1': hp.quniform('sigma_1', low=2.3, high=3.0, q=0.05),
        'eps_1': hp.quniform('eps_1', low=0.04, high=0.45, q=0.01),
        'sigma_2': hp.quniform('sigma_2', low=2.5, high=3.1, q=0.05),
        'eps_2': hp.quniform('eps_2', low=0.16, high=0.39, q=0.01),
        'sigma_3': hp.quniform('sigma_3', low=3.25, high=4.15, q=0.05),
        'eps_3': hp.quniform('eps_3', low=0.04, high=0.318, q=0.002),
        'sigma_4': hp.quniform('sigma_4', low=3.65, high=4.65, q=0.05),
        'eps_4': hp.quniform('eps_4', low=0.142, high=0.606, q=0.002),
        'sigma_nao2': hp.quniform('sigma_nao2', low=2.95, high=3.15, q=0.05),
        'eps_nao2': hp.quniform('eps_nao2', low=0.01, high=0.142, q=0.002),
        'sigma_nao3': hp.quniform('sigma_nao3', low=3.0, high=3.2, q=0.05),
        'eps_nao3': hp.quniform('eps_nao3', low=0.5, high=0.7, q=0.002),
        'sigma_o1ti5': hp.quniform('sigma_o1ti5', low=2.15, high=2.45, q=0.05),
        'eps_o1ti5': hp.quniform('eps_o1ti5', low=0.122, high=0.362, q=0.002),
        'sigma_ho2': hp.quniform('sigma_ho2', low=1.45, high=1.65, q=0.05),
        'eps_ho2': hp.quniform('eps_ho2', low=0.524, high=0.724, q=0.002),
        'sigma_Ho2': hp.quniform('sigma_Ho2', low=2.45, high=2.65, q=0.05),
        'eps_Ho2': hp.quniform('eps_Ho2', low=0.01, high=0.136, q=0.002),
        'sigma_o1o2': hp.quniform('sigma_o1o2', low=3.0, high=3.2, q=0.05),
        'eps_o1o2': hp.quniform('eps_o1o2', low=0.456, high=0.656, q=0.002,),
        'sigma_c1o2': hp.quniform('sigma_c1o2', low=3.12, high=3.32, q=0.05),
        'eps_c1o2': hp.quniform('eps_c1o2', low=0.348, high=0.548, q=0.002)
    }
#    SPACE = {
        #'chg': 0.65,
        #'chg': hp.quniform('chg', low=0.5, high=0.8, q=0.01),
        #'sigma_1': hp.quniform('sigma_1', low=2.2, high=3.5, q=0.05),
        #'eps_1': hp.quniform('eps_1', low=0.04, high=0.45, q=0.01),
        #'sigma_2': hp.quniform('sigma_2', low=2.2, high=3.5, q=0.05),
        #'eps_2': hp.quniform('eps_2', low=0.04, high=0.45, q=0.01),
        #'sigma_3': hp.quniform('sigma_3', low=3.2, high=5.5, q=0.05),
        #'eps_3': hp.quniform('eps_3', low=0.01, high=0.65, q=0.002),
        #'sigma_4': hp.quniform('sigma_4', low=3.2, high=5.5, q=0.05),
        #'eps_4': hp.quniform('eps_4', low=0.01, high=0.65, q=0.002),
        #'sigma_c3o2': hp.quniform('sigma_c3o2', low=2.2, high=8.5, q=0.05),
        #'eps_c3o2': hp.quniform('eps_c3o2', low=0.01, high=0.95, q=0.002),
        #'sigma_cao2': hp.quniform('sigma_cao2', low=2.2, high=8.5, q=0.05),
        #'eps_cao2': hp.quniform('eps_cao2', low=0.01, high=0.95, q=0.002),
        #'sigma_o1ti5': hp.quniform('sigma_o1ti5', low=1.8, high=6.5, q=0.05),
        #'eps_o1ti5': hp.quniform('eps_o1ti5', low=0.01, high=0.95, q=0.002),
        #'sigma_ho2': hp.quniform('sigma_ho2', low=0.9, high=5.5, q=0.05),
        #'eps_ho2': hp.quniform('eps_ho2', low=0.01, high=0.95, q=0.002)
####more precise fit####
        #'sigma_1': hp.quniform('sigma_1', low=2.7, high=3.0, q=0.05),
        #'eps_1': hp.quniform('eps_1', low=0.04, high=0.45, q=0.01),
        #'sigma_2': hp.quniform('sigma_2', low=2.75, high=3.1, q=0.05),
        #'eps_2': hp.quniform('eps_2', low=0.04, high=0.45, q=0.01),
        #'sigma_3': hp.quniform('sigma_3', low=3.4, high=4.5, q=0.05),
        #'eps_3': hp.quniform('eps_3', low=0.01, high=0.65, q=0.002),
        #'sigma_4': hp.quniform('sigma_4', low=3.4, high=4.5, q=0.05),
        #'eps_4': hp.quniform('eps_4', low=0.01, high=0.65, q=0.002),
        #'sigma_nao2': hp.quniform('sigma_nao2', low=1.6, high=4.5, q=0.05),
        #'eps_nao2': hp.quniform('eps_nao2', low=0.01, high=0.95, q=0.002),
        #'sigma_nao3': hp.quniform('sigma_nao3', low=1.6, high=4.5, q=0.05),
        #'eps_nao3': hp.quniform('eps_nao3', low=0.01, high=0.95, q=0.002)
        #'sigma_o1ti5': hp.quniform('sigma_o1ti5', low=1.6, high=5.0, q=0.05),
        #'eps_o1ti5': hp.quniform('eps_o1ti5', low=0.01, high=0.95, q=0.002),
        #'sigma_Ho2': hp.quniform('sigma_Ho2', low=0.9, high=5.5, q=0.05),
        #'eps_Ho2': hp.quniform('eps_Ho2', low=0.01, high=0.95, q=0.002)
        #'sigma_o1o2': hp.quniform('sigma_o1o2', low=1.6, high=5.5, q=0.05),
        #'eps_o1o2': hp.quniform('eps_o1o2', low=0.01, high=0.95, q=0.002)
        #'sigma_c1ti5': hp.quniform('sigma_c1ti5', low=1.5, high=6.5, q=0.05),
        #'eps_c1ti5': hp.quniform('eps_c1ti5', low=0.01, high=0.95, q=0.002)
        #'sigma_c1o2': hp.quniform('sigma_c1o2', low=1.5, high=5.0, q=0.05),
        #'eps_c1o2': hp.quniform('eps_c1o2', low=0.01, high=0.95, q=0.002)
#    }
    seed = int(sys.argv[1])
    best = fmin(objective, SPACE, algo=tpe.suggest, max_evals=200,
         rstate=np.random.RandomState(seed))

    print(best)
    print(E_target)
    #print(MolMM)
    #print(MolQM)
    #print(F_target)
    print(len(Activelll))
    print(Activelll)
    print(Activelli)
    op.close()

if __name__ == '__main__':
    main()
